#!/usr/bin/python

"""
    This file is part of benjap.

    Copyright Graham Jones 2011.

    Benjap is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    banjap is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with benjap (COPYING file).  
    If not, see <http://www.gnu.org/licenses/>.

    $Id$
"""
"""
lunar lander application

"""
import os    

# Start pygame
from appBase import appBase
import pygame
from pygame.locals import *
from random import random, randint, choice
from math import sqrt,sin,cos,pi,atan

class lunarApp(appBase):

    def __init__(self,parent,fullscreen=True):
    
        print "lunarApp.__init__()"
        appBase.__init__(self,parent,fullscreen)

        startHeight = 1e3 # m
        self.timeScale = 1.0  # scale factor to increase time if necessary
        self.scale = 1.1 * startHeight / self.screen_size[1]

        #self.ctrl_x = self.screen_size[0]-100
        #self.ctrl_y = 20
        #print "ctrlxy = (%d,%d)" % (self.ctrl_x,self.ctrl_y)
        self.Sprites = pygame.sprite.RenderUpdates()
        self.origin=(int(self.screen_size[0]/2),
                     int(self.screen_size[1]/2))

        self.screen.fill((0,0,0))
        pygame.display.flip()

        # Load Sounds
        self.engineSound = pygame.mixer.Sound('sounds/rocket_sound.wav')

        # Create Lander
        r = (moon.Rmoon + startHeight)
        v = moon.getOrbitalVelocity(r)
        self.lander = lander(0.,
                             r,
                             v,
                             0.,
                             self,
                             self.Sprites)

        # Create Orbiter
        r = (moon.Rmoon + startHeight)
        v = moon.getOrbitalVelocity(r)
        self.orbiter = lander(-1* r,
                             0.,
                             0.,
                             v,
                             self,
                             self.Sprites)

        #create dashboard
        dbw = 200
        dbh = 60
        dbx = self.screen_size[0]-dbw
        dby = 0
        self.dashboard = dashboard(self,
                                   self.lander,
                                   self.orbiter,
                                   dbx,dby,
                                   dbw,dbh)

        # Create the overview map
        self.overview = overviewMap(self,
                                    self.lander,
                                    self.orbiter,
                                    self.screen_size[0]-dbw,dbh,
                                    dbw,dbw)

    def setOrigin(self,pos):
        self.origin = pos

    def pos2screen(self,pos):
        """ convert a position (x,y) in metres into 
        screen coordinates for the main display.
        """
        x = int(pos[0]/self.scale + self.screen_size[0]/2)
        y = int(pos[1]/self.scale + self.screen_size[1]/2)
        return(x,y)

    def len2screen(self,length):
        """ convert a length in metres into a length in pixels 
        for the main display."""
        return int(length/self.scale)

    def processEvent(self,event):
        if event.type == MOUSEBUTTONDOWN:
            pass
        elif event.type == KEYDOWN:
            if event.key == K_UP:
                #print "K_UP"
                self.lander.incThrust()
                self.engineSound.play()
            elif event.key == K_DOWN:
                #print "K_DOWN"
                self.lander.decThrust()
                self.engineSound.stop()

    def processDisplay(self,delay=0):
        self.dashboard.paint()
        self.overview.paint()

class dashboard():
    """
    A class to display the lander dashboard on the screen
    """
    def __init__(self,parent,orbiter,lander,x,y,w,h):
        self.screen = parent.screen
        self.font18 = parent.font18
        self.lander = lander
        self.orbiter = orbiter
        self.x = x
        self.y = y
        self.w = w
        self.h = h
        self.bgcol = (128,128,128)
        self.lineSpace = 20 # pixels

    def paint(self):
        pygame.draw.rect(self.screen,
                         self.bgcol,
                         (self.x, self.y,
                          self.w, self.h),
                         0)
        self.drawStr(
            "Vert. Vel: %3.1f m/s" % self.lander.dydt,
            1)
        self.drawStr(
            "Hor. Vel: %3.1f m/s" % self.lander.dxdt,
            2)
        self.drawStr(
            "Posn: (%3.0f,%3.0f) m" % (self.lander.x,self.lander.y),
            3)

    def drawStr(self,strVal,lineNo):
        """
        Draw a string onto the dasboard at line no lineNo.
        lines are numbered from 1.
        """
        self.screen.blit(
            self.font18.render(
                strVal, 
                True, 
                (255,255,255)
                ), 
            (self.x,self.y+(lineNo-1)*self.lineSpace)
            )

class overviewMap():
    """
    A class to display the overview map on the screen
    """
    def __init__(self,parent,lander,orbiter,x,y,w,h):
        self.screen = parent.screen
        self.font18 = parent.font18
        self.lander = lander
        self.orbiter = orbiter
        self.x = x
        self.y = y
        self.w = w
        self.h = h
        self.bgcol = (200,200,200)
        self.mooncol = (100,200,100)
        self.landercol = (100,100,200)
        self.orbitercol = (200,100,100)
        self.craftSize = 10
        self.lineSpace = 20 # pixels
        self.scale = self.setScale()

    def setScale(self):
        """ Calculate an appropriate scale for the overview map.
        scale is metres per pixel.
        """
        maxDist = max(self.lander.x,self.lander.y,
                      self.orbiter.x,self.orbiter.y)
        maxDia = 2.*maxDist * 1.1
        self.scale = maxDia / min(self.w,self.h)
        print "overviewMap.setScale - maxDia = %f m, scale=%f m per pixel" % (maxDia,self.scale)
        return self.scale

    def pos2screen(self,pos):
        """ convert a position (x,y) in metres into 
        screen coordinates for the overview map.
        """
        x = self.x + int(pos[0]/self.scale + self.w/2)
        y = self.y + self.h - int(pos[1]/self.scale + self.h/2)  # -ve because zero y at top
        return(x,y)

    def len2screen(self,length):
        """ convert a length in metres into a length in pixels 
        for the overview map."""
        return int(length/self.scale)

    def paint(self):
        pygame.draw.circle(self.screen,
                         self.mooncol,
                         self.pos2screen((0,0)),
                         self.len2screen(moon.Rmoon),
                         0)
        pygame.draw.circle(self.screen,
                         self.landercol,
                         self.pos2screen(self.lander.getPos()),
                         self.craftSize,
                         0)
        pygame.draw.circle(self.screen,
                         self.orbitercol,
                         self.pos2screen(self.orbiter.getPos()),
                         self.craftSize,
                         0)

    def drawStr(self,strVal,lineNo):
        """
        Draw a string onto the dasboard at line no lineNo.
        lines are numbered from 1.
        """
        self.screen.blit(
            self.font18.render(
                strVal, 
                True, 
                (255,255,255)
                ), 
            (self.x,self.y+(lineNo-1)*self.lineSpace)
            )


class craft(pygame.sprite.Sprite):
    def __init__(self,x,y,dxdt,dydt,mainApp,*groups):
        """
        generic representation of a spacecraft - subclass it to customise.
        mainApp is a pointer to the main application to get screen data.
        """
        pygame.sprite.Sprite.__init__(self, *groups)
        self.x = x
        self.dxdt = dxdt
        self.y = y
        self.dydt = dydt
        self.phi = pi/2.  # rotation of craft, pi/2 = vertical

        self.thrust = 0.0
        self.mainApp = mainApp

        self.mass = 10149. # kg (wikipedia Apollo_Lunar_module)
        self.thrust100 = 44400. # N

        image = pygame.image.load("images/lunarLanderOff.jpg")
        self.imageEngineOff = pygame.transform.smoothscale(image,(60,60)).convert_alpha()
        image = pygame.image.load("images/lunarLanderOn.jpg")
        self.imageEngineOn = pygame.transform.smoothscale(image,(60,60)).convert_alpha()
        self.image = self.imageEngineOff

        self.rect = self.image.get_rect(
            center=mainApp.pos2screen(
                (self.x, self.y)
                )
            )
        self.total_time = 0
        self.engineOn = False

    def getPos(self):
        return (self.x,self.y)

    def incThrust(self):
        """ Reduce the engine thrust """
        print "lander.incThrust"
        if not self.engineOn:
            self.engineOn = True
            self.thrust = self.thrust100
            self.image = self.imageEngineOn

    def decThrust(self):
        """ Increase the engine thrust """
        print "lander.decThrust"
        if self.engineOn:
            self.engineOn = False
            self.thrust = 0.0
            self.image = self.imageEngineOff
    
    def update(self, dt_ms):
        """ re-calculate the velocity and actual position after time period dt.
        dt_ms is in miliseconds.
        """
        dt = dt_ms/1000.   # convert to seconds
        self.x = self.x + dt * self.dxdt
        self.y = self.y + dt * self.dydt

        # Force due to gravity
        r2 = self.x*self.x + self.y*self.y  # assumes moon is at 0,0
        if r2 == 0.0:
            r2 = 1e-8
        Fg = moon.G * moon.Mmoon * self.mass / r2  # force acting towards moon.

        # resolve force to x and y components
        if (self.x==0.0):
            theta = pi/2.0
        else:
            theta = atan(self.y/self.x)
        Fg_x = Fg * cos(theta)
        Fg_y = Fg * sin(theta)
        
        # Force due to engine
        if self.engineOn:
            Fe = self.thrust100
        else:
            Fe = 0.

        Fe_x = Fe * cos(self.phi)
        Fe_y = Fe * sin(self.phi)

        # Total force on lander
        Ft_x = Fe_x - Fg_x
        Ft_y = Fe_y - Fg_y

        # vertical acceleration - calculate new velocity
        self.dxdt = self.dxdt + dt * Ft_x/self.mass
        self.dydt = self.dydt + dt * Ft_y/self.mass

        self.rect = self.image.get_rect(
            center=self.mainApp.pos2screen(
                (self.x, self.y)
                )
            )
        print "craft.update: pos=(%f,%f) Ft=(%f,%f), vel=(%f,%f)" \
                   % (self.x,self.y,
                      Ft_x,Ft_y,
                      self.dxdt,self.dydt)

class lander(craft):
    """
    A class to describe the lunar lander and its equations of motion.
    """
    def __init__(self, x,y,dxdt,dydt,mainApp,*groups):
        """
        r and theta are the initial position (r in metres from centre of moon,
        theta in degrees rotation around moon).
        dx and dy are the initial velocity (m/s).
        mainApp is a pointer to the main application to get screen data.
        """
        craft.__init__(self, x,y,dxdt,dydt, mainApp, *groups)
        self.mass = 10149. # kg (wikipedia Apollo_Lunar_module)
        self.thrust100 = 44400. # N

        image = pygame.image.load("images/lunarLanderOff.jpg")
        self.imageEngineOff = \
            pygame.transform.smoothscale(image,
                                         (60,60)).convert_alpha()
        image = pygame.image.load("images/lunarLanderOn.jpg")
        self.imageEngineOn = \
            pygame.transform.smoothscale(image,
                                         (60,60)).convert_alpha()


class moon:
    """
    A simple class to describe the moon
    """
    # Class variables
    Rmoon = 1737.1e3  # m (wikipedia Moon)
    Mmoon = 7.3477e22 # kg (wikipedia Moon)
    G     = 6.67428e-11 # m^3kg^-1s^-2(wikipeia Gravitational_constant)
    
    @staticmethod
    def getOrbitalVelocity(r):
        """
        returns the orbital velocity (m/s) at a distance r from the centre of
        the planet.
        """
        return sqrt(moon.G*moon.Mmoon/r)

if __name__ == "__main__": 
    ml = lunarApp(None,False)
    ml.mainLoop()
